#All the libraries needed to run the scripts are called at the beginning of each script: check the needed libraries there, and install them as recommended on their websites if needed.
#Set working directory to the location of the files. 
#Load libraries
library(Seurat)
options(Seurat.object.assay.version = "v5")
library(ggplot2)
library(patchwork)
library(dplyr)
options(future.globals.maxSize = 1e+09)
library(ggpubr)

#Make helper functions to rotate Seurat Images: Obtained from https://github.com/satijalab/seurat/issues/2702 Amhed Vargas
rotimat=function(foo,rotation){
  if(!is.matrix(foo)){
    cat("Input is not a matrix")
    return(foo)
  }
  if(!(rotation %in% c("180","Hf","Vf"))){
    cat("Rotation should be either L90, R90, 180, Hf or Vf\n")
    return(foo)
  }
  if(rotation == "180"){
    foo <- foo %>% 
      .[, dim(.)[2]:1] %>%
      .[dim(.)[1]:1, ]
  }
  if(rotation == "Hf"){
    foo <- foo %>%
      .[, dim(.)[2]:1]
  }
  
  if(rotation == "Vf"){
    foo <- foo %>%
      .[dim(.)[1]:1, ]
  }
  return(foo)
}

rotateSeuratImage = function(seuratVisumObject, slide = "slice1",rotation="Vf"){
  if(!(rotation %in% c("180","Hf","Vf"))){
    cat("Rotation should be either 180, Hf or Vf\n")
    return(NULL)
  }else{
    seurat.visium = seuratVisumObject
    ori.array = (seurat.visium@images)[[slide]]@image
    img.dim = dim(ori.array)[1:2]/(seurat.visium@images)[[slide]]@scale.factors$lowres
    new.mx <- c()  
    # transform the image array
    for (rgb in 1:3){
      each.mx <- ori.array[,,rgb]
      each.mx.trans <- rotimat(each.mx,rotation)
      new.mx <- c(new.mx, list(each.mx.trans))
    }
    
    # construct new rgb image array
    new.X.dim <- dim(each.mx.trans)[1]
    new.Y.dim <- dim(each.mx.trans)[2]
    new.array <- array(c(new.mx[[1]],
                         new.mx[[2]],
                         new.mx[[3]]), 
                       dim = c(new.X.dim, new.Y.dim, 3))
    
    #swap old image with new image
    seurat.visium@images[[slide]]@image <- new.array
    
    ## step4: change the tissue pixel-spot index
    img.index <- (seurat.visium@images)[[slide]]@coordinates
    
    #swap index
    if(rotation == "Hf"){
      seurat.visium@images[[slide]]@coordinates$imagecol <- img.dim[2]-img.index$imagecol
    }
    
    if(rotation == "Vf"){
      seurat.visium@images[[slide]]@coordinates$imagerow <- img.dim[1]-img.index$imagerow
    }
    
    if(rotation == "180"){
      seurat.visium@images[[slide]]@coordinates$imagerow <- img.dim[1]-img.index$imagerow
      seurat.visium@images[[slide]]@coordinates$imagecol <- img.dim[2]-img.index$imagecol
    }
    
    return(seurat.visium)
  }  
}

#Load Data
load('DESI_ST_Validation_Raw_Data.rdata')

#Remove zero count data
BC_1_DESI_ST <- subset(BC_1_DESI_ST, nCount_Spatial>0)
BC_1_DESI_ST <- subset(BC_1_DESI_ST, nFeature_Spatial>0)
BC_1_ST <- subset(BC_1_ST, nCount_Spatial>0)
BC_1_ST <- subset(BC_1_ST, nFeature_Spatial>0)
LC_1_DESI_ST <- subset(LC_1_DESI_ST, nCount_Spatial>0)
LC_1_DESI_ST <- subset(LC_1_DESI_ST, nFeature_Spatial>0)
LC_1_ST <- subset(LC_1_ST, nCount_Spatial>0)
LC_1_ST <- subset(LC_1_ST, nFeature_Spatial>0)

#Normalize with SCTransform
BC_1_DESI_ST <- SCTransform(BC_1_DESI_ST, assay = "Spatial", verbose = FALSE)
BC_1_ST <- SCTransform(BC_1_ST, assay = "Spatial", verbose = FALSE)
LC_1_DESI_ST <- SCTransform(LC_1_DESI_ST, assay = "Spatial", verbose = FALSE)
LC_1_ST <- SCTransform(LC_1_ST, assay = "Spatial", verbose = FALSE)

#Rotate images to match orientation
BC_1_DESI_ST <- rotateSeuratImage(BC_1_DESI_ST, rotation = "180")
BC_1_DESI_ST <- rotateSeuratImage(BC_1_DESI_ST, rotation = "Vf") 
BC_1_ST <- rotateSeuratImage(BC_1_ST, rotation = "Vf")
LC_1_DESI_ST <- rotateSeuratImage(LC_1_DESI_ST, rotation = "180") 

#Merge Data
BC_1_ST$orig.ident <- gsub("SeuratProject", "ST", BC_1_ST$orig.ident)
BC_1_DESI_ST$orig.ident <- gsub("SeuratProject", "DESI+ST", BC_1_DESI_ST$orig.ident)
head(x = BC_1_ST[[]])
head(x = BC_1_DESI_ST[[]])
LC_1_ST$orig.ident <- gsub("SeuratProject", "ST", LC_1_ST$orig.ident)
LC_1_DESI_ST$orig.ident <- gsub("SeuratProject", "DESI+ST", LC_1_DESI_ST$orig.ident)
head(x = LC_1_ST[[]])
head(x = LC_1_DESI_ST[[]])
BC_1_Merge <- merge(BC_1_ST, BC_1_DESI_ST)
LC_1_Merge <- merge(LC_1_ST, LC_1_DESI_ST)

#Run UMAPS
DefaultAssay(BC_1_Merge) <- "SCT"
VariableFeatures(BC_1_Merge) <- c(VariableFeatures(BC_1_ST), VariableFeatures(BC_1_DESI_ST))
BC_1_Merge <- RunPCA(BC_1_Merge, verbose = FALSE)
BC_1_Merge <- FindNeighbors(BC_1_Merge, dims = 1:30)
BC_1_Merge <- FindClusters(BC_1_Merge, verbose = FALSE)
BC_1_Merge <- RunUMAP(BC_1_Merge, dims = 1:30)
DefaultAssay(LC_1_Merge) <- "SCT"
VariableFeatures(LC_1_Merge) <- c(VariableFeatures(LC_1_ST), VariableFeatures(LC_1_DESI_ST))
LC_1_Merge <- RunPCA(LC_1_Merge, verbose = FALSE)
LC_1_Merge <- FindNeighbors(LC_1_Merge, dims = 1:30)
LC_1_Merge <- FindClusters(LC_1_Merge, verbose = FALSE)
LC_1_Merge <- RunUMAP(LC_1_Merge, dims = 1:30)

#Plot UMAPS
BC_Ident <- DimPlot(BC_1_Merge, pt.size = 1.3, reduction = "umap", group.by = "ident", shuffle = TRUE) + 
  ggtitle("Cluster Identity") + theme(plot.title = element_text(hjust = 0.5))
BC_OrigIdent <- DimPlot(BC_1_Merge, pt.size = 1.3, reduction = "umap", group.by = "orig.ident", shuffle = TRUE) + 
  ggtitle("Sample Identity") + theme(plot.title = element_text(hjust = 0.5))
wrap_plots(BC_Ident, BC_OrigIdent) 

LC_Ident <- DimPlot(LC_1_Merge, pt.size = 1.3, reduction = "umap", group.by = "ident", shuffle = TRUE) + ggtitle("Cluster Identity") +
  theme(plot.title = element_text(hjust = 0.5))
LC_OrigIdent <- DimPlot(LC_1_Merge, pt.size = 1.3, reduction = "umap", group.by = "orig.ident", shuffle = TRUE) + ggtitle("Sample Identity") +
  theme(plot.title = element_text(hjust = 0.5))
wrap_plots(LC_Ident, LC_OrigIdent) 
 
SpatialDimPlot(BC_1_Merge, pt.size.factor = 1.3, label = TRUE, label.size = 2.5, image.alpha = 0, crop = FALSE)
SpatialDimPlot(LC_1_Merge, pt.size.factor = 1.3, label = TRUE, label.size = 2.5, crop = FALSE, image.alpha = 0)

#Violin plots: Genes/Spot and UMI Counts/Spot
BC_nCount_Spatial <- VlnPlot(object = BC_1_Merge, group.by = 'orig.ident', 
  features = 'nCount_Spatial', y.max = 150000, cols = c("#f8766d", "#00bfc4"), pt.size = 0) + 
  ggtitle("BC-1 UMI Counts/Spot") + xlab("") + ylab("UMI Counts") + NoLegend() + 
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)
BC_nFeature_Spatial <- VlnPlot(object = BC_1_Merge, group.by = 'orig.ident', 
  features = 'nFeature_Spatial', y.max = 11500, cols = c("#f8766d", "#00bfc4"), pt.size = 0) + 
  ggtitle("BC-1 Genes/Spot") + xlab("") + ylab("Gene Counts") + NoLegend() + 
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)
LC_nCount_Spatial <- VlnPlot(object = LC_1_Merge, group.by = 'orig.ident', 
  features = 'nCount_Spatial', y.max = 150000, cols = c("#f8766d", "#00bfc4"), pt.size = 0) + 
  ggtitle("LC-1 UMI Counts/Spot") + xlab("") + ylab("UMI Counts") + NoLegend() + 
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)
LC_nFeature_Spatial <- VlnPlot(object = LC_1_Merge, group.by = 'orig.ident',
  features = 'nFeature_Spatial', y.max = 11500, cols = c("#f8766d", "#00bfc4"), pt.size = 0) + 
  ggtitle("LC-1 Genes/Spot") + xlab("") + ylab("Gene Counts") + NoLegend() + 
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)
wrap_plots(BC_nCount_Spatial, LC_nCount_Spatial)
wrap_plots(BC_nFeature_Spatial, LC_nFeature_Spatial)

#Correlation scatter plots UMI Counts/Gene
averaged_expression_BC <- (AggregateExpression(BC_1_Merge, group.by = "orig.ident")) 
averaged_expression_LC <- (AggregateExpression(LC_1_Merge, group.by = "orig.ident"))

head(averaged_expression_BC$Spatial)
head(averaged_expression_LC$Spatial)

ave_ex_spat_BC <- as.data.frame(averaged_expression_BC$Spatial)
colnames(ave_ex_spat_BC) <- c('DESI_ST', 'ST')
head(ave_ex_spat_BC)

ave_ex_spat_LC <- as.data.frame(averaged_expression_LC$Spatial)
colnames(ave_ex_spat_LC) <- c('DESI_ST', 'ST')
head(ave_ex_spat_LC)

BC_Scat <- ggscatter(ave_ex_spat_BC, x = "DESI_ST", y = "ST", add = "reg.line") +
  stat_cor(label.x = 250000, label.y = 150000, digits = 4) + ggtitle("BC-1 UMI Counts/Gene") +
  xlab("DESI+ST") + ylab("ST") + NoLegend() + ylim(0, 700000) + xlim(0, 650000) +
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)
LC_Scat <- ggscatter(ave_ex_spat_LC, x = "DESI_ST", y = "ST", add = "reg.line") +
  stat_cor(label.x = 250000, label.y = 150000, , digits = 4) + ggtitle("LC-1 UMI Counts/Gene") +
  xlab("DESI+ST") + ylab("ST") + NoLegend() + ylim(0, 700000) + xlim(0, 650000) +
  theme(plot.title = element_text(hjust = 0.5)) + font("title", size = 14, face = "bold") +
  font("xlab", size = 12) + font("ylab", size = 12)

wrap_plots(BC_Scat, LC_Scat)
